package com.lht.arithmetic.tree.binaryTree.avl;

/**
 * 先实现一个二叉树，然后在考虑保证他平衡
 * 然后写一个测试类，检查平衡的稳定性
 *
 * @author Eric
 * @version 1.0
 * @date 2019-11-21
 */
public class MyAvlTree<T extends Comparable<T>> {


    private AvlNode<T> root;


    public boolean contains(T element) {
        AvlNode<T> node = root;
        while (node != null) {
            if (node.element.compareTo(element) == 0) {
                return true;
            } else if (node.element.compareTo(element) > 0) {
                node = node.left;
            } else if (node.element.compareTo(element) < 0) {
                node = node.right;
            }
        }
        return false;
    }

    public T search(T element) {
        AvlNode<T> node = root;
        while (node != null) {
            if (node.element.compareTo(element) == 0) {
                return node.element;
            } else if (node.element.compareTo(element) > 0) {
                node = node.left;
            } else if (node.element.compareTo(element) < 0) {
                node = node.right;
            }
        }
        return null;
    }

    public boolean add(T element) {
        if (root == null) {
            root = new AvlNode<T>(element, null, null);
            return true;
        }
        AvlNode<T> node = root;
        while (node != null) {
            if (node.element.compareTo(element) == 0) {
                return true;
            } else if (node.element.compareTo(element) > 0) {
                if (node.left == null) {
                    node.left = new AvlNode<T>(element, null, null);
                    return true;
                } else {
                    node = node.left;
                }
            } else if (node.element.compareTo(element) < 0) {
                if (node.right == null) {
                    node.right = new AvlNode<T>(element, null, null);
                    return true;
                } else {
                    node = node.right;
                }
            }

        }
        return false;
    }

    public boolean remove(T element) {
        AvlNode<T> node = root;
        while (node != null) {
            if (node.element.compareTo(element) == 0) {
                AvlNode<T> tmp = node.left;
                if (tmp == null) {
                    tmp = node.right;
                }
                return true;
            } else if (node.element.compareTo(element) > 0) {
                node = node.left;
            } else if (node.element.compareTo(element) < 0) {
                node = node.right;
            }
        }
        return false;
    }

    public static class AvlNode<T extends Comparable<T>> {
        private T element;

        private AvlNode<T> left;

        private AvlNode<T> right;

        public AvlNode() {
            this(null, null, null);
        }

        public AvlNode(T element, AvlNode<T> left, AvlNode<T> right) {
            this.element = element;
            this.left = left;
            this.right = right;
        }

        public T getElement() {
            return element;
        }

        public void setElement(T element) {
            this.element = element;
        }

        public AvlNode<T> getLeft() {
            return left;
        }

        public void setLeft(AvlNode<T> left) {
            this.left = left;
        }

        public AvlNode<T> getRight() {
            return right;
        }

        public void setRight(AvlNode<T> right) {
            this.right = right;
        }
    }

    public AvlNode<T> getRoot() {
        return root;
    }

    public void setRoot(AvlNode<T> root) {
        this.root = root;
    }
}
